﻿#pragma once

#include <iostream>
#include <cstdlib>
using namespace std;

namespace tree
{
	enum Colour
	{
		RED,
		BLACK
	};

	template<class K, class V>
	struct RBTreeNode
	{
		// 这里更新控制平衡也要加入parent指针
		pair<K, V> _kv;
		RBTreeNode<K, V>* _left;
		RBTreeNode<K, V>* _right;
		RBTreeNode<K, V>* _parent;
		Colour _col;

		RBTreeNode(const pair<K, V>& kv)
			:_kv(kv)
			, _left(nullptr)
			, _right(nullptr)
			, _parent(nullptr)
		{}
	};

	template<class K, class V>
	class Red_Black_Tree
	{
	public:
		typedef RBTreeNode<K, V> Node;

		Red_Black_Tree() = default;

		~Red_Black_Tree()
		{}

		void RotateR(Node* parent)
		{
			Node* SubL = parent->_left;
			Node* SubLR = SubL->_right;
			Node* pParent = parent->_parent;

			parent->_left =  SubLR;
			if (SubLR)
				SubLR->_parent = parent;

			SubL->_right = parent;

			if (pParent == nullptr)
			{
				_root == SubL;
				_root->_parent = nullptr;
			}
			else
			{
				if (pParent->_left == parent)
				{
					pParent->_left = SubL;
				}
				else
				{
					pParent->_right = SubL;
				}
				SubL->_parent = nullptr;
			}
		}

		void RotateL(Node* parent)
		{
			Node* SubR = parent->_right;
			Node* SubRL = SubR->_left;
			Node* pParent = parent->_parent;

			parent->_right = SubRL;
			if (SubRL)
				SubRL->_parent = parent;

			SubR->_left = parent;
			if (pParent == nullptr)
			{
				_root = SubR;
				SubR->_parent = nullptr;
			}
			else
			{
				if (pParent->_left == parent)
				{
					pParent->_left = SubR;
				}
				else
				{
					pParent->_right = SubR;
				}
				SubR->_parent = pParent;
			}
		}


		bool insert(const pair<K, V> kv)
		{
			if (_root == nullptr)
			{
				_root = new Node(kv);
				_root->_col = BLACK;

				return true;
			}

			Node* cur = _root;
			Node* parent = nullptr;

			while (cur)
			{
				if (kv.first > cur->_kv.first)
				{
					parent = cur;
					cur = cur->_right;
				}
				else if (kv.first < cur->_kv.first)
				{
					parent = cur;
					cur = cur->_left;
				}
				else
				{
					return false;
				}
			}

			cur = new Node(kv);
			cur->_parent = parent;
			if (kv.first > parent->_kv.first)
				parent->_right = cur;
			else
				parent->_left = cur;

			while (parent && parent->_col == RED)
			{
				Node* grandfather = parent->_parent;
				if (grandfather->_left == parent)
				{
					Node* Uncle = grandfather->_right;
					if (Uncle && Uncle->_col == RED)
					{
						Uncle->_col = parent->_col = BLACK;
						grandfather->_col = RED;
						//cur = grandfather;
						//parent = cur->_parent;
					}
					else
					{
						if (parent->_left == cur)
						{
							RotateR(grandfather);
							parent->_col = BLACK;
							grandfather->_col = RED;
						}
						else
						{
							RotateL(parent);
							RotateR(grandfather);
							cur->_col = BLACK;
							grandfather->_col = RED;
						}
						break;
					}
				}
				else
				{
					Node* Uncle = grandfather->_left;
					if (Uncle && Uncle->_col == RED)
					{
						Uncle->_col = parent->_col = BLACK; 
						grandfather->_col = RED;
						//cur = grandfather;
						//parent = cur->_parent;
					}
					else
					{
						if (parent->_right == cur)
						{
							RotateL(grandfather);
							parent->_col = BLACK;
							grandfather->_col = RED;
						}
						else
						{
							RotateR(parent);
							RotateL(grandfather);
							cur->_col = BLACK;
							grandfather->_col = RED;
						}
						break;
					}

					cur = grandfather;
					parent = cur->_parent;
				}
			}


			return true;
		}

		bool Check(Node* root, int blackNum, const int refNum)
		{
			if (root == nullptr)
			{
				// 前序遍历⾛到空时，意味着⼀条路径⾛完了 
				//cout << blackNum << endl;

				if (refNum != blackNum)
				{
					cout << "存在⿊⾊结点的数量不相等的路径" << endl;
					return false;
				}
				return true;
			}

			// 检查孩⼦不太⽅便，因为孩⼦有两个，且不⼀定存在，反过来检查⽗亲就⽅便多了 
			if (root->_col == RED && root->_parent->_col == RED)
			{
				cout << root->_kv.first << "存在连续的红⾊结点" << endl;
				return false;
			}
			if (root->_col == BLACK)
			{
				blackNum++;
			}
			return Check(root->_left, blackNum, refNum)
				&& Check(root->_right, blackNum, refNum);
		}


		bool IsBalance()
		{
			if (_root == nullptr)
				return true;
			if (_root->_col == RED)
				return false;

			// 参考值 
			int refNum = 0;
			Node* cur = _root;
			while (cur)
			{
				if (cur->_col == BLACK)
				{
					++refNum;
				}
				cur = cur->_left;
			}
			return Check(_root, 0, refNum);
		}


	private:
		Node* _root = nullptr;
	};

	void test()
	{
		Red_Black_Tree<int, int> t;
		srand(time(NULL));

		int n = 100000;

		clock_t start = clock();
		while (n--)
		{
			int i = n + rand();
			t.insert({ i , i });
		}
		clock_t end = clock();
		cout << end - start << endl;
		cout << t.IsBalance() << endl;
		//cout << t.Check() << endl;
	}
}